
from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Optional
import uuid
import os
import aiofiles

from scorpio.shared.ingest import document_processor
from scorpio.shared.generation import generation_service
from scorpio.core.providers import db_manager

from scorpio.core.common import get_logger
logger = get_logger(__name__)

app = FastAPI(title="RAG System API", version="1.0.0")

class QueryRequest(BaseModel):
    question: str
    top_k: Optional[int] = 5
    max_tokens: Optional[int] = 1000
    temperature: Optional[float] = 0.7

class QueryResponse(BaseModel):
    question: str
    answer: str
    retrieved_context: List[dict]
    context_count: int

class DocumentUploadResponse(BaseModel):
    document_id: int
    external_document_id: str
    title: str
    total_chunks: int
    status: str

@app.on_event("startup")
async def startup_event():
    """应用启动时初始化服务"""
    await db_manager.initialize()

@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭时清理资源"""
    await db_manager.close()

@app.post("/documents/upload", response_model=DocumentUploadResponse)
async def upload_document(
    file: UploadFile = File(...),
    title: Optional[str] = None,
    chunk_size: int = 1000,
    chunk_overlap: int = 200
):
    """上传并处理文档"""
    
    # 创建临时文件
    file_ext = os.path.splitext(file.filename)[1].lower()
    temp_filename = f"temp_{uuid.uuid4()}{file_ext}"
    
    try:
        # 保存上传的文件
        async with aiofiles.open(temp_filename, 'wb') as temp_file:
            content = await file.read()
            await temp_file.write(content)
        
        # 处理文档
        result = await document_processor.process_document(
            file_path=temp_filename,
            document_id=str(uuid.uuid4()),
            title=title or file.filename,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap
        )
        
        return result
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}")
    
    finally:
        # 清理临时文件
        if os.path.exists(temp_filename):
            os.remove(temp_filename)

@app.post("/query", response_model=QueryResponse)
async def query_rag_system(request: QueryRequest):
    """查询 RAG 系统"""
    try:
        result = await generation_service.generate_answer(
            question=request.question,
            top_k=request.top_k,
            max_tokens=request.max_tokens,
            temperature=request.temperature
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")

@app.get("/query/stream")
async def query_rag_system_stream(
    question: str,
    top_k: int = Query(5, ge=1, le=20),
    max_tokens: int = Query(1000, ge=100, le=4000),
    temperature: float = Query(0.7, ge=0.0, le=1.0)
):
    """流式查询 RAG 系统"""
    logger.info(f"Received query: {question}")
    async def generate():
        async for chunk in generation_service.generate_answer_stream(
            question=question,
            top_k=top_k,
            max_tokens=max_tokens,
            temperature=temperature
        ):
            yield chunk
    
    return StreamingResponse(generate(), media_type="text/plain")

@app.get("/documents/{document_id}/stats")
async def get_document_stats(document_id: int):
    """获取文档统计信息"""
    stats = await db_manager.get_document_stats(document_id)
    if not stats:
        raise HTTPException(status_code=404, detail="文档不存在")
    return stats


@app.get("/rag/health")
async def health_check():
    """健康检查"""
    logger.info("Health check requested")
    return {"status": "healthy", "service": "RAG System"}